import torch


__all__ = ["Attention"]


import torch
import torch.nn as nn
import torch.nn.functional as F


class Attention(nn.Module):
    def __init__(self, enc_hid_dim, dec_hid_dim):
        super().__init__()
        self.attn = nn.Linear(enc_hid_dim + dec_hid_dim, dec_hid_dim)
        self.v = nn.Linear(dec_hid_dim, 1, bias=False)

    def forward(self, hidden, encoder_outputs, mask):
        # hidden: [batch size, dec hid dim], encoder_outputs: [batch size, src len, enc hid dim]
        src_len = encoder_outputs.shape[1]

        hidden = hidden.unsqueeze(1).repeat(
            1, src_len, 1
        )  # Repeat decoder hidden state src_len times
        energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim=2)))
        attention = self.v(energy).squeeze(2)  # [batch size, src len]

        attention = attention.masked_fill(mask == 0, -1e10)  # Apply mask
        return F.softmax(attention, dim=1)
